use pnet::datalink::{self, NetworkInterface};
use pnet::datalink::Channel::Ethernet;
use pnet::packet::ethernet::{EthernetPacket, EtherTypes};
use pnet::packet::ip::IpNextHeaderProtocols;
use pnet::packet::ipv4::Ipv4Packet;
use pnet::packet::ipv6::Ipv6Packet;
use pnet::packet::tcp::TcpPacket;
use pnet::packet::udp::UdpPacket;
use pnet::packet::Packet;

fn main() {
    // 获取所有可用的网络接口
    let interfaces = datalink::interfaces();

    // 选择要监控的网络接口，这里简单选择第一个可用的非回环接口
    let interface = match find_default_interface(&interfaces) {
        Some(iface) => iface,
        None => {
            eprintln!("No suitable network interface found.");
            return;
        }
    };

    println!("Monitoring interface: {}", interface.name);

    // 创建一个数据包捕获通道
    let (_, mut rx) = match datalink::channel(&interface, Default::default()) {
        Ok(Ethernet(tx, rx)) => (tx, rx),
        Ok(_) => {
            eprintln!("Unhandled channel type.");
            return;
        }
        Err(e) => {
            eprintln!("Error creating datalink channel: {}", e);
            return;
        }
    };

    loop {
        // 从通道中接收数据包
        match rx.next() {
            Ok(packet) => {
                // 解析以太网帧
                if let Some(eth_packet) = EthernetPacket::new(packet) {
                    match eth_packet.get_ethertype() {
                        // 处理 IPv4 数据包
                        EtherTypes::Ipv4 => {
                            if let Some(ipv4_packet) = Ipv4Packet::new(eth_packet.payload()) {
                                handle_ipv4_packet(&ipv4_packet);
                            }
                        }
                        // 处理 IPv6 数据包
                        EtherTypes::Ipv6 => {
                            if let Some(ipv6_packet) = Ipv6Packet::new(eth_packet.payload()) {
                                handle_ipv6_packet(&ipv6_packet);
                            }
                        }
                        _ => {}
                    }
                }
            }
            Err(e) => {
                eprintln!("Error receiving packet: {}", e);
            }
        }
    }
}

// 查找默认网络接口的辅助函数
fn find_default_interface(interfaces: &[NetworkInterface]) -> Option<NetworkInterface> {
    for interface in interfaces {
        if interface.is_up() && !interface.is_loopback() {
            return Some(interface.clone());
        }
    }
    None
}

// 处理 IPv4 数据包
fn handle_ipv4_packet(ipv4_packet: &Ipv4Packet) {
    println!("IPv4 Packet  Source IP: {}  Destination IP: {}  Protocol: {:?}", ipv4_packet.get_source(), ipv4_packet.get_destination(), ipv4_packet.get_next_level_protocol());

    match ipv4_packet.get_next_level_protocol() {
        // 处理 TCP 数据包
        IpNextHeaderProtocols::Tcp => {
            if let Some(tcp_packet) = TcpPacket::new(ipv4_packet.payload()) {
                println!("  TCP Packet     Source Port: {}    Destination Port: {}    Payload Length: {}", tcp_packet.get_source(), tcp_packet.get_destination(), tcp_packet.payload().len());
                // 获取 TCP 负载（body 内容）
                let tcp_body = tcp_packet.payload();

                // 尝试将负载转换为字符串
                match std::str::from_utf8(tcp_body) {
                    Ok(s) => println!("TCP body as string: {}", s),
                    Err(_) => println!("TCP body is not valid UTF-8"),
                }
            }
        }
        // 处理 UDP 数据包
        IpNextHeaderProtocols::Udp => {
            if let Some(udp_packet) = UdpPacket::new(ipv4_packet.payload()) {
                println!("  UDP Packet    Source Port: {}    Destination Port: {}    Payload Length: {}", udp_packet.get_source(), udp_packet.get_destination(), udp_packet.payload().len());
            }
        }
        _ => {}
    }
}

// 处理 IPv6 数据包
fn handle_ipv6_packet(ipv6_packet: &Ipv6Packet) {
    println!("IPv6 Packet  Source IP: {}  Destination IP: {}  Next Header: {:?}", ipv6_packet.get_source(), ipv6_packet.get_destination(), ipv6_packet.get_next_header());

    match ipv6_packet.get_next_header() {
        // 处理 TCP 数据包
        IpNextHeaderProtocols::Tcp => {
            if let Some(tcp_packet) = TcpPacket::new(ipv6_packet.payload()) {
                println!("  TCP Packet    Source Port: {}    Destination Port: {}    Payload Length: {}", tcp_packet.get_source(), tcp_packet.get_destination(), tcp_packet.payload().len());
                let tcp_body = tcp_packet.payload();

                // 尝试将负载转换为字符串
                match std::str::from_utf8(tcp_body) {
                    Ok(s) => println!("TCP body as string: {}", s),
                    Err(_) => println!("TCP body is not valid UTF-8"),
                }

            }
        }
        // 处理 UDP 数据包
        IpNextHeaderProtocols::Udp => {
            if let Some(udp_packet) = UdpPacket::new(ipv6_packet.payload()) {
                println!("  UDP Packet    Source Port: {}    Destination Port: {}    Payload Length: {}", udp_packet.get_source(), udp_packet.get_destination(), udp_packet.payload().len());
            }
        }
        _ => {}
    }
}